"""\
This is ImagProcPhysTools - a collection of small convenience tools to make your life 
easier and focus on what really matters.

Version 0.2 - 06.11.2011 - Fixed Qt plotting problems. 
Version 0.3 - 14.11.2011 - Added entervalue function
Version 0.4 - 02.12.2011 - Added qindex for Fourier-space indices that match with FFTs
Version 0.5 - 12.12.2011 - Added pause and prop
"""

import numpy as np
import Image as PIL
import matplotlib
import sys
import time

# Fix tif import problem
PIL._MODE_CONV['I;16'] = (PIL._ENDIAN + 'u2', None)

# Improved interactive behavior or matplotlib 
import threading
if matplotlib.get_backend().lower().startswith('qt4'):
    mpl_backend = 'qt'
    from PyQt4 import QtGui
    gui_yield_call = QtGui.qApp.processEvents
elif matplotlib.get_backend().lower().startswith('wx'):
    mpl_backend = 'wx'
    import wx
    gui_yield_call = wx.Yield
elif matplotlib.get_backend().lower().startswith('gtk'):
    mpl_backend = 'gtk'
    import gtk
    def gui_yield_call():
        gtk.gdk.threads_enter()
        while gtk.events_pending():
            gtk.main_iteration(True)
        gtk.gdk.flush()
        gtk.gdk.threads_leave()
else:
    mpl_backend = None


__all__ = ['hitreturn', 'pause', 'askchoice', 'imread', 'entervalue', 'qindex',\
           'prop']

if mpl_backend is not None:
    class _Pause(threading.Thread):
        def __init__(self, timeout, message):
            self.message = message
            self.timeout = timeout
            self.ct = True
            threading.Thread.__init__(self)
        def run(self):
            sys.stdout.flush()
            if self.timeout < 0:
                raw_input(self.message)
            else:
                if self.message is not None:
                    print self.message
                time.sleep(self.timeout)
            self.ct = False
    def pause(timeout=-1, message=None):
        """\
        Pause the execution of a script while leaving matplotlib figures responsive.
        By default, execution is resumed only after hitting return. If timeout >= 0,
        the execution is resumed after timeout seconds.
        """
        if message is None:
            if timeout < 0:
                message = 'Paused. Hit return to continue.'
        h = _Pause(timeout, message)
        h.start()
        while h.ct:
            gui_yield_call()
            time.sleep(.01)
else:
    def pause(timeout=-1, message=None):
        """\
        Pause the execution of a script.
        By default, execution is resumed only after hitting return. If timeout >= 0,
        the execution is resumed after timeout seconds.
        This version of pause is not GUI-aware (this happens if the matplotlib
        backend is not supported).
        """
        if timeout < 0:
            if message is None:
                message = 'Paused. Hit return to continue.'
            raw_input(message)
        else:
            if message is not None:
                print message
            time.sleep(timeout)

# For backward compatibility. "pause" should always be used now.
hitreturn = pause

def askchoice(choices, default=0, prompt='Please make your choice'):
    """\
    Ask the user to pick an element in a list or a dictionary
    """
    if len(choices) > 15:
        raise RuntimeError('askchoice: Too many elements in the list.')
    print prompt + ':'
    if isinstance(choices, dict):
        for i,kv in enumerate(choices.items()):
            print '%3d) %s' % (i,kv[0])
    else:
        for i,v in enumerate(choices):
            print '%3d) %s' % (i,v)
    choice = None
    while choice is None:
        r = raw_input('[default: %d] >>> ' % default)
        if not r:
            return default
        try:
            r = int(r)
        except:
            print 'Invalid choice'
            continue
        if r < 0 or r >= len(choices):
            print 'Invalid choice'
            continue
        return r

def entervalue(prompt='Please enter some value', default = 1):
    """\
    Ask the user to enter some value
    """
    r = None
    while r is None:
        r = raw_input(prompt + '[default: %d] >>> ' % default)
        if not r:
            return default
        try:
            r = float(r)
        except:
            print 'Invalid choice'
            continue
        return r


def imread(filename):
    """\
    Read an image file and returns a numpy array and metadeta attached to the file.
    """
    im = PIL.open(filename)
    data = np.array(im)
    meta = dict(im.info.items())
    meta['format'] = im.format
    meta['format_description'] = im.format_description
    meta['mode'] = im.mode
    return data, meta

def qindex(shape):
    """\
    This function returns the Fourier-space indices that match with FFTs. 
    """
    n0,n1 = shape
    x0,x1 = np.indices(shape)
    q0 = (x0 + n0//2) % n0 - n0//2
    q1 = (x1 + n1//2) % n1 - n1//2
    return (q0,q1)

def prop(a, lam, z):
    """\
    Near field free-space propagation of array a over a distance z, with wavelength
    lam (pixel units).
    """
    sh = a.shape
    k = 2.*np.pi/lam
    q0,q1 = qindex(sh)
    dq0, dq1 = 2*np.pi/sh[0], 2*np.pi/sh[1]
    q_orth2 = ((q0*dq0)**2 + (q1*dq1)**2)
    near_field_propagator = np.exp(-1j * z  * q_orth2 / 2. /k )
    return np.fft.ifftn(np.fft.fftn(a) * near_field_propagator)

def test():
    """\
    Testing python's plotting behavior for various contexts.
    This should behave well and consistently wether run from the command
    line, within ipython or within spyder. It should also work ideally
    in Windows, Linux and MacOS.
    """
    from matplotlib import pyplot as plt
    plt.ion()

    plt.figure(1)
    plt.plot([0,1], [0,1])
    plt.title('This is a graph. Try to zoom and confirm that it is responsive.')
    hitreturn("""Execution blocked, but GUI should still be responsive. 
Hit return to continue""")   
    plt.close(1)

if __name__ == "__main__":
    test()
